"""Training the network"""

from argparse import ArgumentParser

import numpy as np

config = dict()

# config for inputs
config['input_shape']         = (96, 96, 80)
config['batch_size']          = 4

# config for network
config['network']             = 'mynetwork'
config['weight_decay']        = False

# config for training
config['iterations']          = 20000
config['learning_rate']       = 1e-3
config['learning_rate_decay'] = False


def train(output_path, shape, platform):
    """Train unet

    Args:
        output_path: path of saved model
        shape: input image shape
        platform: deep learning platform used
    """

    from Utils import DataLoader
    from Model import Unet, Trainer

    data_loader = DataLoader("./Data", shape=shape, platform=platform)
    print("Data Loader initialize complete.")

    net = Unet(config)
    trainer = Trainer(config, net)
    path = trainer.train(data_loader, output_path)

if __name__ == "__main__":

    parser = ArgumentParser()

    parser.add_argument('-bs', '--batchsize', type=int, default=2,
                        help="Batch size")
    parser.add_argument('-it', '--iterations', type=int, default=50000,
                        help="Max iterations")
    parser.add_argument('-p', '--path', type=str, default="./Checkpoint",
                        help="Output path")
    parser.add_argument('-s', '--shape', nargs='+', type=int,
                        default=[96, 96, 80], help="Image shape")
    parser.add_argument('-f', '--platform', type=str, default="tensorflow",
                        help="Deep learning platform used")
    args = parser.parse_args()

    assert len(args.shape) == 3, "Image dimension should be 3."

    if args.platform.lower() != "tensorflow":
        raise NotImplementedError("%s not implemented!" % args.platform)

    config['input_shape'] = args.shape
    config['batch_size']  = args.batchsize
    config['iterations']  = args.iterations

    train(args.path, args.shape, args.platform)
